import torch
import torch.nn as nn
import torch.nn.functional as F

Device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(Device)


class Embedding(nn.Embedding):
    def __init__(self, vocab_size, embedding_dim):
        super(Embedding, self).__init__(vocab_size, embedding_dim,device=Device)


if __name__ == '__main__':
    embedding = Embedding(10, 5)
    embedding = embedding.to(device=Device)

    data = torch.randint(10, size=(5, 2))
    data = data.to(device=Device)
    print(data)
    print(data.dtype)
    # print(embedding.weight.shape)
    # print(embedding(data))
    print(embedding(data).shape)
    print(embedding(data))
    print(embedding.weight.device)
    print(data.device)
    print('end')
